/*
 * @Author: wanli wanliofficial@aliyun.com
 * @Date: 2024-09-02 13:18:22
 * @LastEditors: wanli wanliofficial@aliyun.com
 * @LastEditTime: 2025-05-13 17:46:11
 * @FilePath: \locust4j\src\main\java\com\github\myzhan\locust4j\rpc\ZeromqClient.java
 * @Description: 
 * 
 * Copyright (c) 2025 by wanli email: wanliofficial@aliyun.com, All Rights Reserved. 
 */
package com.github.myzhan.locust4j.rpc;

import java.io.IOException;

import com.github.myzhan.locust4j.message.Message;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.zeromq.ZMQ;
import org.zeromq.ZContext;
import org.zeromq.ZMQException;
import org.zeromq.SocketType;

/**
 * Locust used to support both plain-socket and zeromq. Since Locust 0.8, it
 * removes the plain-socket implementation.
 * Locust4j only supports zeromq.
 *
 * @author myzhan
 */
public class ZeromqClient implements Client {

    private static final Logger logger = LoggerFactory.getLogger(ZeromqClient.class);

    private final ZContext context = new ZContext();
    private final String identity;
    private final ZMQ.Socket dealerSocket;

    public ZeromqClient(String host, int port, String nodeID) {
        this.identity = nodeID;

        // 创建一个 Socket 对象，指定类型为 REQ（请求）
        this.dealerSocket = context.createSocket(SocketType.DEALER);

        // 连接到服务器。这里假设服务器在本地运行，监听于端口 5555。
        this.dealerSocket.setIdentity(this.identity.getBytes());
        boolean connected = this.dealerSocket.connect(String.format("tcp://%s:%d", host, port));
        if (connected) {
            logger.debug("Locust4j is connected to master({}:{})", host, port);
        } else {
            logger.debug("Locust4j isn't connected to master({}:{}), please check your network situation", host, port);
        }
    }

    @Override
    public Message recv() throws IOException {
        try {
            byte[] bytes = this.dealerSocket.recv();
            return new Message(bytes);
        } catch (ZMQException ex) {
            throw new IOException("Failed to receive ZeroMQ message", ex);
        }
    }

    @Override
    public void send(Message message) throws IOException {
        byte[] bytes = message.getBytes();
        this.dealerSocket.send(bytes);
    }

    @Override
    public void close() {
        dealerSocket.close();
        context.close();
    }
}
